#include <device_launch_parameters.h>

#include <cuda_test_utils.cuh>

// 向量加法核函数，grid和block都是1D
__global__ void add(float* x, float* y, float* z, int n)
{
    // 获取全局索引
    int index = threadIdx.x + blockIdx.x * blockDim.x;
    // 步长
    int stride = blockDim.x * gridDim.x;
    for (int i = index; i < n; i += stride)
    {
        z[i] = x[i] + y[i];
    }
}

// 测试类
class Te1_1_Add : public CudaTest
{
public:
    void run() override
    {
        int N = 1 << 20;
        int nBytes = N * sizeof(float);
        // 分配主机内存
        float *x, *y, *z;
        x = (float*)malloc(nBytes);
        y = (float*)malloc(nBytes);
        z = (float*)malloc(nBytes);
        // 初始化数据
        for (int i = 0; i < N; ++i)
        {
            x[i] = 10.0;
            y[i] = 20.0;
        }

        // 分配设备内存
        float *d_x, *d_y, *d_z;
        cudaMalloc((void**)&d_x, nBytes);
        cudaMalloc((void**)&d_y, nBytes);
        cudaMalloc((void**)&d_z, nBytes);

        // 从主机内存复制数据到设备内存
        cudaMemcpy((void*)d_x, (void*)x, nBytes, cudaMemcpyHostToDevice);
        cudaMemcpy((void*)d_y, (void*)y, nBytes, cudaMemcpyHostToDevice);
        // 定义核函数执行配置
        dim3 blockSize(256);
        dim3 gridSize((N + blockSize.x - 1) / blockSize.x);
        // 执行核函数
        add<<<gridSize, blockSize>>>(d_x, d_y, d_z, N);
        // 等待所有线程执行完毕
        cudaDeviceSynchronize();
        // 从设备内存复制数据到主机内存
        cudaMemcpy((void*)z, (void*)d_z, nBytes, cudaMemcpyDeviceToHost);
        // 检查执行结果
        float maxError = 0.0;
        for (int i = 0; i < N; i++) maxError = fmax(maxError, fabs(z[i] - 30.0));
        MYINFO("Maximum error: {0}", maxError);

        // 释放设备内存
        cudaFree(d_x);
        cudaFree(d_y);
        cudaFree(d_z);
        // 释放主机内存
        free(x);
        free(y);
        free(z);

        // return 0;
    }

    std::string name() const override { return "te_add"; }
};

REGISTER_TEST(Te1_1_Add);